open Base
open Unique_id.Ids

type vars = Src_id.t Var_id.map
type deps = Src_id.set
type 'a samp = Src.values -> 'a

type +'a t =
  { eval : vars -> deps * 'a samp
  } [@@ocaml.unboxed]

let of_src (s : Src_id.t) : Dyn.t t =
  let deps = Set.singleton (module Src_id) s in
  let sample m = Map.find_exn m s in
  { eval = fun _ -> (deps, sample) }

let map {eval} ~f =
  let eval vs =
    let (deps, sample) = eval vs in
    (deps, fun m -> f (sample m))
  in {eval}

let both a b =
  let eval vs =
    let (a_deps, a_sample) = a.eval vs in
    let (b_deps, b_sample) = b.eval vs in
    (Set.union a_deps b_deps,
     fun m -> (a_sample m, b_sample m))
  in {eval}

(**********)

type 'a var =
  { id : Var_id.t
  ; ty : 'a Dyn.rtti }

let create_var () =
  { id = Var_id.create ()
  ; ty = Dyn.create_reified () }

let var_id {id;_} = id

let upcast var x = var.ty#to_dyn x
let upcast1 var f dyn =
  match var.ty#of_dyn dyn with
  | None -> dyn
  | Some(x) -> var.ty#to_dyn (f x)

let of_var var =
  let eval vs =
    match Map.find vs var.id with
    | None -> failwith "variable not in scope"
    | Some(src_id) ->
       let deps = Set.singleton (module Src_id) src_id in
       let sample m =
         match var.ty#of_dyn (Map.find_exn m src_id) with
         | None -> failwith "source has incorrect type"
         | Some(v) -> v
       in
       (deps, sample)
  in
  {eval}

let empty_vars : vars =
  Var_id.empty_map

let bind vars ~var ~src_id =
  Map.set vars ~key:var.id ~data:src_id
